/******************************************************************************
 * Copyright 2022 The Airos Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/

#include "air_service/modules/perception-fusion/pipeline/multi_sensor_fusion_component.h"

#include <iostream>

#include "air_middleware_rate.h"
#include "base/common/log.h"
#include "base/common/time_util.h"
#include "air_service/modules/perception-fusion/algorithm/air_fusion/air_object_fusion.h"
#include "air_service/modules/perception-fusion/pipeline/tools/msf_config_manager.h"
#include "air_service/modules/perception-fusion/pipeline/tools/msf_gflags.h"
#include "air_service/modules/perception-fusion/pipeline/tools/msf_trans_tools.h"
#include "yaml-cpp/yaml.h"

namespace airos {
namespace perception {
namespace msf {

using airos::base::Singleton;

bool SENSOR_FUSION_COMPONENT::Init() {
  if (InitConfig() == false) {
    LOG_ERROR << "InitConfig failed.";
    return false;
  }

  if (InitFusion() == false) {
    LOG_ERROR << "InitFusion failed.";
    return false;
  }

  if (InitNode() == false) {
    LOG_ERROR << "InitNode failed.";
    return false;
  }

  if (InitInputList() == false) {
    LOG_ERROR << "InitList failed.";
    return false;
  }

#ifdef MSF_VIZ
  InitVisualization();
#endif
  
  main_loop_.reset(
      new std::thread(std::bind(&SENSOR_FUSION_COMPONENT::MainLoop, this)));

  initialized_ = true;
  return true;
}

bool SENSOR_FUSION_COMPONENT::InitConfig() {
  auto& msf_config =
      Singleton<tools::MsfConfigManager>::get_instance()->GetMsfConfig();
  base_msf_param_.channel_num = msf_config.input_size();
  base_msf_param_.time_diff = msf_config.timestamp_param().time_diff();
  base_msf_param_.missing_time_threshold =
      msf_config.timestamp_param().missing_time_threshold();
  base_msf_param_.reverse_time_threshold =
      msf_config.timestamp_param().reverse_time_threshold();

  for (int i = 0; i < msf_config.timestamp_param().reference_seq_size(); ++i) {
    base_msf_param_.reference_seq.insert(
        msf_config.timestamp_param().reference_seq(i));
  }
  base_msf_param_.fusion_params_file = tools::FLAGS_fusion_params_file;
  return true;
}

bool SENSOR_FUSION_COMPONENT::InitFusion() {
  air_fusion_ = std::make_shared<msf::AirObjectFusion>();  // TODO fgq
  if (air_fusion_->Init(base_msf_param_) == false) {
    LOG_ERROR << "fusion init failed.";
    return false;
  }
  return true;
}

bool SENSOR_FUSION_COMPONENT::InitNode() {
  auto& msf_config =
      Singleton<tools::MsfConfigManager>::get_instance()->GetMsfConfig();
  const std::string& node_name = msf_config.node_name();
  node_ = std::make_shared<middleware::AirMiddlewareNode>(node_name);

  if (node_ == nullptr) {
    return false;
  }

  for (int i = 0; i < msf_config.input_size(); ++i) {
    auto reader = node_->CreateReader<PerceptionObstacles>(
        msf_config.input(i).channel_name(),
        std::bind(&SENSOR_FUSION_COMPONENT::InternalCallback, this, i,
                  std::placeholders::_1));
    if (reader == nullptr) {
      return false;
    }
    reader_list_.push_back(reader);
  }

  const std::string& perception_channel = msf_config.perception_channel();
  perception_writer_ =
      node_->CreateWriter<PerceptionObstacles>(perception_channel);
  if (perception_writer_ == nullptr) {
    return false;
  }
  return true;
}

bool SENSOR_FUSION_COMPONENT::InitInputList() {
  auto& msf_config =
      Singleton<tools::MsfConfigManager>::get_instance()->GetMsfConfig();
  fusion_input_list_.resize(base_msf_param_.channel_num);

  for (size_t i = 0; i < base_msf_param_.channel_num; ++i) {
    auto& input = msf_config.input(i);

    fusion_input_list_[i].index = i;
    fusion_input_list_[i].input_type =
        static_cast<msf::InputType>(input.input_type());
    fusion_input_list_[i].use_header_time = input.use_header_time();
    fusion_input_list_[i].main_sensor = input.main_sensor();
    fusion_input_list_[i].collision_line_fusion = input.collision_line_fusion();
  }

  return true;
}

#ifdef MSF_VIZ
bool SENSOR_FUSION_COMPONENT::InitVisualization() {
  LOG_INFO << "Enable visualization.";
  msf_visual_kernal_ = std::make_shared<msf::MSFVisualKernal>();
  auto fusion_param =
      Singleton<tools::MsfConfigManager>::get_instance()->GetFusionParams();
  double chisquare_critic = fusion_param.score_params().chi_square_critical();
  msf_visual_kernal_->Init(tools::FLAGS_viz_mode, chisquare_critic);
  return true;
}
#endif

void SENSOR_FUSION_COMPONENT::InternalCallback(
    int sequence,
    const std::shared_ptr<const PerceptionObstacles>& callback_msg) {
  if (!initialized_ || is_exit_) {
    return;
  }

  {
    auto& fusion_input = fusion_input_list_[sequence];
    std::unique_lock<std::mutex> lck(loop_mutex_);
    if (callback_msg->error_code() == PerceptionErrorCode::ERROR_NONE) {
      tools::TransTools::PbsToObjects(*callback_msg, &fusion_input);
    }
  }
  new_msg_.notify_one();
}

void SENSOR_FUSION_COMPONENT::MainLoop() {
  double output_frequency = Singleton<tools::MsfConfigManager>::get_instance()
                                ->GetMsfConfig()
                                .output_frequency();
  airos::middleware::AirMiddlewareRate rate(output_frequency);
  while (!is_exit_) {
    std::vector<msf::FusionInput> fusion_input_list;
    {
      std::unique_lock<std::mutex> lck(loop_mutex_);
      new_msg_.wait(lck);
      fusion_input_list = fusion_input_list_;
    }
    double start_time = airos::base::TimeUtil::GetCurrentTime();
    MsfProc(fusion_input_list);
    LOG_INFO << "MsfProc() time delay: "
             << airos::base::TimeUtil::GetCurrentTime() - start_time
             << "us, with " << fusion_output_.objects.size() << " objects";
    rate.Sleep();
  }
}

void SENSOR_FUSION_COMPONENT::MsfProc(
    std::vector<msf::FusionInput>& fusion_input_list) {
  msf::FusionOutput fusion_output;
  int error_code = air_fusion_->Process(fusion_input_list, fusion_output);
  auto output_msg = std::make_shared<PerceptionObstacles>();
#ifdef MSF_VIZ
  msf_visual_kernal_->UpdateObstacles(fusion_input_list, fusion_output);
#endif

  SerializeMsg(fusion_output, output_msg, error_code);
  if (CheckTimestamp(output_msg)) {  // TODO fgq
    perception_writer_->Write(output_msg);
  }

  seq_num_++;
}

void SENSOR_FUSION_COMPONENT::SerializeMsg(
    const msf::FusionOutput& message,
    std::shared_ptr<PerceptionObstacles> output_msg, int error_code) {
  tools::TransTools::ObjectsToPbs(message, output_msg.get());
  airos::header::Header* header = output_msg->mutable_header();
  header->set_timestamp_sec(airos::base::TimeUtil::GetCurrentTime());
  header->set_module_name("v2x_fusion");
  header->set_sequence_num(seq_num_);

  if (error_code == static_cast<int>(msf::base::Error::NONE)) {
    long timestamp = long(air_fusion_->Timestamp() * 1e9);
    header->set_camera_timestamp(timestamp);
    output_msg->set_error_code(
        airos::perception::PerceptionErrorCode::ERROR_NONE);
  } else {
    output_msg->set_error_code(
        airos::perception::PerceptionErrorCode::ERROR_PROCESS);
  }
}

bool SENSOR_FUSION_COMPONENT::CheckTimestamp(
    std::shared_ptr<PerceptionObstacles> output_msg) {
  auto header = output_msg->header();

  static uint64_t pre_camera_timestamp = 0;

  if (header.camera_timestamp() <= pre_camera_timestamp) {
    LOG_ERROR << std::setprecision(16) << "Sensor timestamp error: "
              << header.camera_timestamp() - pre_camera_timestamp
              << ", pre: " << pre_camera_timestamp
              << ", cur: " << header.camera_timestamp();
    return false;
  }

  pre_camera_timestamp = header.camera_timestamp();
  return true;
}

SENSOR_FUSION_COMPONENT::~SENSOR_FUSION_COMPONENT() {
  is_exit_ = true;
  new_msg_.notify_one();
  if (main_loop_->joinable()) {
    main_loop_->join();
  }
}

}  // namespace msf
}  // namespace perception
}  // namespace airos